import numpy as np
from functools import reduce


class Lagrange_Polynomial_Interpolation(object):
    """Lagrange_Polynomial_Interpolation 

    base function: Ln(x)=\sum_{k=0}^{n} yk l_k(x)
    Author: Junno
    Date: 2022-03-01
    """

    def __init__(self, x, y):
        """__init__ 

        Args:
            x ([np.array]): [x of points]
            y ([np.array]): [y of points]
        """
        self.N = x.shape[0]
        self.x = np.copy(x)
        self.y = np.copy(y)

    def calculate_lk(self, x, k,):
        """calculate_lk(x)

        base function: lk(x)=∏——{j=0,j!=k}^{n} (x-xj)/(xk-xj)
        Args:
            x ([float]): [x]
            k ([int]): [k]

        Returns:
            [float]: [lk(x)]
        """
        ta = reduce(lambda x, y: x*y, [x-self.x[i]
                    for i in range(self.N) if i != k])
        tb = reduce(lambda x, y: x*y, [self.x[k]-self.x[i]
                    for i in range(self.N) if i != k])
        return ta/tb

    def predict(self, x):
        """predict 

        Args:
            x ([float]): [x]

        Returns:
            [float]: [predict y]
        """
        return np.sum([self.y[k]*self.calculate_lk(x, k) for k in range(self.N)])


class mini_I(object):
    """
    class for calculating the interpolation of order k+1 from two component of order k
    base funtion: ans=y0+(y1-y0)/(x1-x0)*(x-x0)

    Author: Junno
    Date: 2022-03-01
    """

    def __init__(self, x0=0, x1=0, y0=0, y1=0, first=False):
        """__init__ 
        Args:
            x0 (float, optional): [component of the base function]. Defaults to 0.
            x1 (float, optional): [component of the base function]. Defaults to 0.
            y0 (float, optional): [component of the base function]. Defaults to 0.
            y1 (float, optional): [component of the base function]. Defaults to 0.
            first (bool, optional): [whether the first element of each order list]. Defaults to False.
        """
        self.x0 = x0
        self.x1 = x1
        self.y0 = y0
        self.y1 = y1
        self.first = first
        # record answer for fast calculation
        self.last_answer = {}

    def __call__(self, x=0.):
        """__call__ 

        Args:
            x ([float]): [interpolation node]

        Returns:
            [answer]: [recursive call or the recorded values]

        """
        if x not in self.last_answer.keys():
            # here x0,x1,y0,y1 can be a mini_I class to call for input x
            answer = self.y0(x)+(self.y1(x)-self.y0(x))/(self.x1 -
                                                         self.x0)*(x-self.x0) if not self.first else self.y0
            # record
            self.last_answer[x] = answer
            return answer
        else:
            if x in self.last_answer.keys():
                return self.last_answer[x]


class Successive_Linear_Interpolation_2D(object):
    """Successive_Linear_Interpolation_2D 逐次线性插值

    Author: Junno
    Date: 2022-03-01
    """

    def __init__(self, x, y, base_order=0):
        """__init__ 

        Args:
            x ([np.array]): [x of points]
            y ([np.array]): [y of points]
            base_order ([int]): [interpolation order]
        """
        self.N = x.shape[0]
        assert base_order < self.N, 'The base_order should be less than the number of points'
        self.x = np.copy(x)
        self.y = np.copy(y)
        self.base_order = 0
        # flag of calculated
        self.calculated = False

        # ref:《数值分析》 李庆扬、王能超、易大义（华中科大出版社）
        self.F = lambda i, j: mini_I(
            self.x[i-j], self.x[i], self.I_set[i-1][j-1], self.I_set[i][j-1], first=False)

    def calculate_I(self, x, eps=1e-6, order=1, check=True):
        """calculate_I calculate y for a given x and order with threshold eps

        Args:
            x ([float]): [x]
            eps ([float], optional): [threshold to stop process]. Defaults to 1e-6.
            order ([int], optional): [interpolation order]. Defaults to 1.
        """
        assert order < self.N, 'The order should be less than the number of points'

        if not self.calculated:

            self.base_order = self.N if order is None else order

            self.I_set = [[]]

            self.I_set[0].append(mini_I(y0=self.y[0], first=True))

            last = 0

            for i in range(1, self.base_order+1):
                self.I_set += [[]]
                self.I_set[i].append(mini_I(y0=self.y[i], first=True))
                for k in range(1, i+1):
                    self.I_set[i].append(self.F(i, k))

                # early stop when the difference between of two results from k order and k+1 order is less than eps
                if i > 1:
                    new = self.I_set[i][-1](x)
                    # print(i, abs(last-new) <= eps)
                    if abs(last-new) <= eps:
                        if i < (self.base_order+1)-1:
                            if check:
                                print(
                                    "Meet with absotion condition of eps={} and the current order of interpolation is {}".format(eps, i))
                            self.base_order = i
#                             self.I_set.pop(-1)

                        self.calculated = True
                        return new

                last = self.I_set[i][-1](x)

                # print(i,)
                # self.print(x)

            self.calculated = True

            return self.I_set[self.base_order][self.base_order](x)

        else:
            last = self.I_set[0][0](x)
            for i in range(1, order+1):

                if i > self.base_order:
                    self.I_set += [[]]
                    self.I_set[i].append(mini_I(y0=self.y[i], first=True))
                    for k in range(1, i+1):
                        self.I_set[i].append(self.F(i, k))

                new = self.I_set[i][-1](x)
                if abs(last-new) <= eps:
                    if check:
                        print(
                            "Meet with absotion condition of eps={} and the current order of interpolation is {}".format(eps, i))
                    return new
                last = new

            return self.I_set[order][order](x)

    def add_point(self, x, y):
        """add_point 

        Args:
            x ([np.array]): [x of points]
            y ([np.array]): [y of points]
        """
        M = x.shape[0]
        self.x = np.concatenate((self.x, x), axis=0)
        self.y = np.concatenate((self.y, y), axis=0)
        self.N = self.N+M

    def print(self, x):
        """print show I_set for a given x

        Args:
            x ([float]): [x]
        """
        for i in range(len(self.I_set)):
            print('I_{}:'.format(i), end='\t')
            for j in range(len(self.I_set[i])):
                print(self.I_set[i][j](x), end='\t')
            print('\n')


class Equidistant_Newton_Linear_Interpolation:
    """Equidistant_Newton_Linear_Interpolation 等距节点牛顿插值法

    Author: Junno
    Date: 2022-03-09
    """
    def __init__(self,x,y,h):
        """__init__ 

        Args:
            x ([np.array]): [x of points]
            y ([np.array]): [y of points]
            h ([float]): [gap between of points]
        """
        self.N=y.shape[0]
        self.base_order=0
        self.eps=eps
        self.x=x
        self.y=y
        self.h=h
        
        # function to calculate binomial coefficients
        self.binomial_coefficient=lambda j,k:reduce(lambda x,y:x*y,[j-i for i in range(k)])/reduce(lambda x,y:x*y, [i for i in range(1,k+1)]) if k>0 else 1
        
        # forward difference matrix
        self.forward_delta_matrix=self.cal_forward_(self.y,self.N)
                    
        # backward difference matrix
        self.backward_nabla_matrix=self.cal_backward_(self.y,self.N)

        if self.base_order == 0:
            self.base_order=self.N-1
    
    def cal_forward_(self,y,N):
        """cal_forward_ 

        Args:
            y ([np.array]): [y of points]
            N ([int]): [number of points]

        Returns:
            [np.darray]: [forward_delta_matrix]
        """        
        temp=np.zeros((N,N))
        temp[:,0]=y
        for j in range(1,N):
            for i in range(N-j):
                temp[i][j]=temp[i+1][j-1]-temp[i][j-1]
        return temp 
        
    def cal_backward_(self,y,N):
        """cal_backward_ 

        Args:
            y ([np.array]): [y of points]
            N ([int]): [number of points]

        Returns:
            [np.darray]: [backward_nabla_matrix]
        """
        temp=np.zeros((N,N))
        temp[:,0]=y
        for j in range(1,N):
            for i in range(j,N):
                temp[i][j]=temp[i-1][j-1]-temp[i][j-1]
        return temp 
        
    def calulate(self,x,eps=1e-3,direction=None,show_error=True):
        """__init__ 

        Args:
            x ([float]): [x to calculate y]
            eps ([float], optional): [threshold to stop process]. Defaults to 1e-6.
            direction ([str]): [difference direction: 'forward' or 'backward']
            show_error ([bool], optional): [whether to error]. Defaults to True.
            
        Return:
            y ([float]): [predict y value]
        """
        if direction==None:
        # half to be forward and half to be backward
            direction='forward' if x<self.x[0]+self.N/2*self.h else 'backward'
        # find nearest points index
        ind=np.argmin(abs(self.x-x))
        
        if direction == 'forward':
            t=(x-self.x[ind])/self.h
            
            count=np.sum(np.abs(self.forward_delta_matrix[ind])>eps)
            ans= sum([self.binomial_coefficient(t,k)*self.forward_delta_matrix[ind][k] for k in range(count)])
            
            if show_error:
                ny=np.sort(np.concatenate((self.y,[ans]),axis=0), axis=0)
                nm=self.cal_forward_(ny,self.N+1)
                error=np.abs(self.binomial_coefficient(t,count)*np.max(np.abs(nm[:,count])))
                print("MAX INTERPOLATION ERROR:{}".format(error))
            return ans
            
        else:
            t=-(x-self.x[ind])/self.h       
            count=np.sum(np.abs(self.backward_nabla_matrix[ind])>eps)
            ans=sum([(-1)**k*self.binomial_coefficient(-t,k)*self.backward_nabla_matrix[ind][k] for k in range(count)])
            if show_error:
                ny=np.sort(np.concatenate((self.y,[ans]),axis=0), axis=0)
                nm=self.cal_backward_(ny,self.N+1)
                error=np.abs(self.binomial_coefficient(-t,count)*np.max(np.abs(nm[:,count])))
                print("MAX INTERPOLATION ERROR:{}".format(error))
                
            return ans
    
    def add_point(self, x, y):
        """add_point 

        Args:
            x ([np.array]): [x of points]
            y ([np.array]): [y of points]
        """
        M = x.shape[0]
        self.x = np.concatenate((self.x, x), axis=0)
        self.y = np.concatenate((self.y, y), axis=0)
        self.N = self.N+M
                